import argparse
import time
from pathlib import Path
import smtplib
from email.mime.base import MIMEBase
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random
import face_recognition
import os
import cv2

from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, \
    strip_optimizer, set_logging, increment_path
from utils.plots import plot_one_box
from utils.torch_utils import select_device, load_classifier, time_synchronized

# 发送邮件的代码
# coding:utf-8
import smtplib
from email.mime.text import MIMEText
import time

from email.mime.multipart import MIMEMultipart
from email.header import Header
from email import encoders
from email.mime.base import MIMEBase
from email.utils import parseaddr, formataddr

from email.mime.image import MIMEImage


def SendMail(sender, receivers, cc_mail, mail_pass, content, file1, file2):
    # 第三方 SMTP 服务
    mail_host = "smtp.qq.com"  # 设置服务器

    # message = MIMEText(content, 'plain', 'utf-8')#正文内容   plain代表纯文本

    # 构造一个MIMEMultipart对象代表邮件本身
    message = MIMEMultipart()
    message.attach(MIMEText(content, 'html', 'utf-8'))  # 正文内容   plain代表纯文本,html代表支持html文本

    message['From'] = sender
    message['To'] = ','.join(receivers)  # 与真正的收件人的邮箱不是一回事
    message['Cc'] = ','.join(cc_mail)

    subject = '军容风纪检查结果-%s' % time.ctime()
    message['Subject'] = subject  # 邮件标题

    # 添加文件到附件
    with open(file1, 'rb') as f:
        # MIMEBase表示附件的对象
        mime = MIMEBase('text', 'txt', filename=file1)
        # filename是显示附件名字
        mime.add_header('Content-Disposition', 'attachment', filename=file1)
        # 获取附件内容
        mime.set_payload(f.read())
        encoders.encode_base64(mime)
        # 作为附件添加到邮件
        message.attach(mime)

    with open(file2, 'rb') as f:
        # MIMEBase表示附件的对象
        mime = MIMEBase('text', 'txt', filename=file2)
        # filename是显示附件名字
        mime.add_header('Content-Disposition', 'attachment', filename=file2)
        # 获取附件内容
        mime.set_payload(f.read())
        encoders.encode_base64(mime)
        # 作为附件添加到邮件
        message.attach(mime)

    try:
        smtpObj = smtplib.SMTP_SSL(mail_host, 465)
        smtpObj.login(sender, mail_pass)
        smtpObj.sendmail(sender, receivers + cc_mail, str(message))  # message.as_string()
        smtpObj.quit()
        print("邮件发送成功")

    except smtplib.SMTPException:
        print("发送成功")


##发送邮件代码结束

def detect(save_img=False):
    source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
    webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
        ('rtsp://', 'rtmp://', 'http://'))

    # Directories
    save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))  # increment run
    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

    # Initialize
    set_logging()
    device = select_device(opt.device)
    half = device.type != 'cpu'  # half precision only supported on CUDA

    # Load model
    model = attempt_load(weights, map_location=device)  # load FP32 model
    imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size
    if half:
        model.half()  # to FP16

    # Second-stage classifier
    classify = False

    if classify:
        modelc = load_classifier(name='resnet101', n=2)  # initialize
        modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()

    # Set Dataloader
    vid_path, vid_writer = None, None

    if webcam:
        view_img = True
        cudnn.benchmark = True  # set True to speed up constant image size inference
        dataset = LoadStreams(source, img_size=imgsz)
    else:
        save_img = True
        dataset = LoadImages(source, img_size=imgsz)

    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]

    # Run inference
    t0 = time.time()
    img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
    _ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once

    for path, img, im0s, vid_cap in dataset:
        img = torch.from_numpy(img).to(device)
        img = img.half() if half else img.float()  # uint8 to fp16/32
        img /= 255.0  # 0 - 255 to 0.0 - 1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)

        # Inference
        t1 = time_synchronized()
        pred = model(img, augment=opt.augment)[0]

        # Apply NMS
        pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
        t2 = time_synchronized()

        # Apply Classifier
        if classify:
            pred = apply_classifier(pred, modelc, img, im0s)

        # Process detections
        for i, det in enumerate(pred):  # detections per image
            if webcam:  # batch_size >= 1
                p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
            else:
                p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)

            p = Path(p)  # to Path
            save_path = str(save_dir / p.name)  # img.jpg
            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # img.txt
            s += '%gx%g ' % img.shape[2:]  # print string
            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
            if len(det):
                # Rescale boxes from img_size to im0 size
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
                # Print results
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += f'{n} {names[int(c)]}s, '  # add to string

                array_of_img = []

                # Write results
                for *xyxy, conf, cls in reversed(det):
                    if save_txt:  # Write to file
                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                        line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh)  # label format
                        with open(txt_path + '.txt', 'a') as f:
                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or view_img:  # Add bbox to image
                        label = f'{names[int(cls)]} {conf:.2f}'
                        plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)

            # Print time (inference + subject = "test report"NMS)
            print(f'{s}Done. ({t2 - t1:.3f}s)')
            # Stream results
            if view_img:
                cv2.imshow(str(p), im0)
                if cv2.waitKey(1) == ord('q'):  # q to quit
                    raise StopIteration
            # Save results (image with detections)
            if save_img:
                if dataset.mode == 'image':
                    cv2.imwrite(save_path, im0)
                else:  # 'video'
                    if vid_path != save_path:  # new video
                        vid_path = save_path
                        if isinstance(vid_writer, cv2.VideoWriter):
                            vid_writer.release()  # release previous video writer

                        fourcc = 'mp4v'  # output video codec
                        fps = vid_cap.get(cv2.CAP_PROP_FPS)
                        w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                        h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                        vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
                    vid_writer.write(im0)

            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

            f = open('/home/c4isr/yolov5/final_result.txt', mode='w+')
            name = compare_faces("/home/c4isr/人脸识别项目/photos")
            f.write(name + '\n')
            if len(det):
                j = 0
                for i in det[:, -1].unique():
                    if int(i) == 0:
                        f.write('该同志存在领带下拉问题\n')
                    if int(i) == 2:
                        f.write('该同志已领佩戴领带\n')
                    if int(i) == 3:
                        f.write('该同志存在佩戴耳环问题\n')
                    if int(i) == 1:
                        j = j + 1
                    continue
                if j == 6:
                    f.write('该同志常服纽扣已扣好\n')
                if j > 0 and j < 6:
                    f.write('该同志常服纽扣未扣好\n')
                continue
            continue
            f.close()

    if save_txt or save_img:
        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
        print(f"Results saved to {save_dir}{s}")

    print(f'Done. ({time.time() - t0:.3f}s)')


def compare_faces(directory_name):
    for filename in os.listdir(directory_name):
        path = directory_name + "/" + filename
        image1 = face_recognition.load_image_file(directory_name + "/" + filename)
        face_locations1 = face_recognition.face_locations(image1)

        face_landmarks_list1 = face_recognition.face_landmarks(image1)
        list_of_face_encodings1 = face_recognition.face_encodings(image1)

        print(face_locations1)
        print(face_landmarks_list1)
        print(list_of_face_encodings1)
        image2 = face_recognition.load_image_file("/home/c4isr/人脸识别项目/test.png")
        face_locations2 = face_recognition.face_locations(image2)
        face_landmarks_list2 = face_recognition.face_landmarks(image2)
        list_of_face_encodings2 = face_recognition.face_encodings(image2)
        print(face_locations2)
        print(face_landmarks_list2)
        print(list_of_face_encodings2)

        for face_encoding2 in list_of_face_encodings2:
            matches = face_recognition.compare_faces(list_of_face_encodings1, face_encoding2, 0.6)

        print(matches)

    return os.path.splitext(os.path.basename(path))[0]


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='data/images', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    opt = parser.parse_args()
    print(opt)
    array_of_img = []
    with torch.no_grad():
        if opt.update:  # update all models (to fix SourceChangeWarning)
            for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
                detect()
                strip_optimizer(opt.weights)


        else:
            detect()

    ##发送邮件给领导
    sender = '895143767@qq.com'  # 用户名与发送方
    receivers = ['895143767@qq.com']  # 接收邮件，可设置为你的QQ邮箱或者其他邮箱
    cc_mail = ['895143767@qq.com']  # 抄送人

    # 口令,QQ邮箱是输入授权码，在qq邮箱设置 里用验证过的手机发送短信获得，g不含空格
    mail_pass = "uqonygiuqxbvbeia"

    content = '''
            <h1>一名同志军容风纪检查结果汇报和相关材料。</h1>
          '''
    file1 = '/home/c4isr/yolov5/final_result.txt'
    file2 ='/home/c4isr/yolov5/runs/detect/exp/1.jpg'
    SendMail(sender, receivers, cc_mail, mail_pass, content, file1, file2)


